import pygame
import math
import time

# 游戏常量
SCREEN_WIDTH = 1200
SCREEN_HEIGHT = 600
GROUND_HEIGHT = 500
BALL_RADIUS = 8
GRAVITY = 980
AIR_FRICTION = 0.1
MAX_POWER = 800
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
RESTITUTION = 0.6  # 反弹系数
FRICTION = 0.8  # 地面摩擦力

pygame.init()
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
font = pygame.font.Font(None, 24)


class ThrowRecord:
    def __init__(self, distance, duration, angle, power):
        self.distance = distance
        self.duration = duration
        self.angle = angle
        self.power = power


class GameState:
    def __init__(self):
        self.reset()
        self.records = []
        self.enable_roll = True  # 是否开启滚动效果

    def reset(self):
        self.angle = 45
        self.power = 300
        self.ball_pos = [50, GROUND_HEIGHT - BALL_RADIUS]
        self.ball_vel = [0, 0]
        self.flying = False
        self.rolling = False
        self.start_time = 0
        self.land_time = 0
        self.current_distance = 0
        self.current_time = 0
        self.first_impact = True  # 首次落地标记


state = GameState()


def draw_ground():
    pygame.draw.line(screen, WHITE, (0, GROUND_HEIGHT),
                     (SCREEN_WIDTH, GROUND_HEIGHT), 2)


def draw_arrow():
    start = state.ball_pos
    radians = math.radians(state.angle)
    length = state.power / 4
    end = (
        start[0] + length * math.cos(radians),
        start[1] - length * math.sin(radians)
    )
    pygame.draw.line(screen, RED, start, end, 3)
    arrow_size = 15
    points = [
        (end[0] - arrow_size * math.cos(radians - math.pi / 6),
         end[1] + arrow_size * math.sin(radians - math.pi / 6)),
        end,
        (end[0] - arrow_size * math.cos(radians + math.pi / 6),
         end[1] + arrow_size * math.sin(radians + math.pi / 6))
    ]
    pygame.draw.lines(screen, RED, False, points, 3)


def draw_ui():
    # Control panel
    texts = [
        f"ANG: {state.angle}° (<- ->)",
        f"PWR: {int(state.power)}N (^ v)",
        "[SPACE] Shoot"
    ]
    for i, text in enumerate(texts):
        surf = font.render(text, True, WHITE)
        screen.blit(surf, (20, 20 + i * 25))

    # Flight data
    if state.flying or state.rolling:
        texts = [
            f"TIME: {state.current_time:.2f}s",
            f"DIST: {state.current_distance:.2f}m"
        ]
        for i, text in enumerate(texts):
            surf = font.render(text, True, WHITE)
            screen.blit(surf, (20, 120 + i * 25))

    # Leaderboard
    screen.blit(font.render("TOP 10 RECORDS:", True, WHITE), (20, 220))
    headers = ["RANK", "DIST", "TIME", "ANG", "PWR"]
    screen.blit(font.render("  ".join(headers), True, WHITE), (20, 250))

    for i, record in enumerate(state.records[:10]):
        text = f"{i + 1:2d}  {record.distance:6.2f}m  {record.duration:.2f}s  {record.angle:.1f}°  {record.power:.0f}N"
        screen.blit(font.render(text, True, WHITE), (20, 280 + i * 25))

    # Roll physics checkbox
    pygame.draw.rect(screen, WHITE, (SCREEN_WIDTH - 150, 20, 130, 30), 2)
    check_text = "Roll Physics: " + ("ON" if state.enable_roll else "OFF")
    screen.blit(font.render(check_text, True, WHITE), (SCREEN_WIDTH - 140, 25))
    if state.enable_roll:
        pygame.draw.line(screen, WHITE, (SCREEN_WIDTH - 152, 22), (SCREEN_WIDTH - 142, 30), 2)
        pygame.draw.line(screen, WHITE, (SCREEN_WIDTH - 142, 30), (SCREEN_WIDTH - 126, 18), 2)


def update_physics(dt):
    if state.flying or state.rolling:
        # 垂直运动
        state.ball_vel[1] += GRAVITY * dt
        state.ball_vel[1] *= math.exp(-AIR_FRICTION * dt)

        # 水平运动
        state.ball_vel[0] *= math.exp(-AIR_FRICTION * dt)
        if not state.flying:  # 滚动时增加地面摩擦
            state.ball_vel[0] *= math.exp(-FRICTION * dt)

        # 更新位置
        state.ball_pos[0] += state.ball_vel[0] * dt
        state.ball_pos[1] += state.ball_vel[1] * dt

        # 落地检测
        if state.ball_pos[1] + BALL_RADIUS >= GROUND_HEIGHT:
            if state.first_impact:
                # 首次落地记录数据
                state.land_time = time.time()
                state.records.append(ThrowRecord(
                    state.current_distance,
                    state.current_time,
                    state.angle,
                    state.power
                ))
                state.records.sort(key=lambda x: -x.distance)
                state.records = state.records[:10]
                state.first_impact = False

                if state.enable_roll:
                    # 进入滚动状态
                    state.flying = False
                    state.rolling = True
                    state.ball_vel[1] *= -RESTITUTION  # 反弹
                    state.ball_vel[0] *= FRICTION  # 水平减速
                    state.ball_pos[1] = GROUND_HEIGHT - BALL_RADIUS
                else:
                    state.reset()

            else:
                # 后续弹跳
                state.ball_vel[1] *= -RESTITUTION
                state.ball_vel[0] *= FRICTION
                state.ball_pos[1] = GROUND_HEIGHT - BALL_RADIUS

        # 滚动超时检测
        if state.rolling and (time.time() - state.land_time > 3):
            state.reset()

        # 更新当前数据
        state.current_time = time.time() - state.start_time
        state.current_distance = (state.ball_pos[0] - 50) / 10


def handle_input():
    mouse_pos = pygame.mouse.get_pos()

    # 处理复选框点击
    if pygame.mouse.get_pressed()[0]:
        if SCREEN_WIDTH - 150 < mouse_pos[0] < SCREEN_WIDTH - 20 and 20 < mouse_pos[1] < 50:
            state.enable_roll = not state.enable_roll
            time.sleep(0.2)  # 防连点

    keys = pygame.key.get_pressed()

    if keys[pygame.K_r]:  # 重置游戏
        state.reset()
        state.records = []

    if not state.flying and not state.rolling:
        # 调整角度
        if keys[pygame.K_LEFT]:
            state.angle = max(0, state.angle - 1)
        if keys[pygame.K_RIGHT]:
            state.angle = min(89, state.angle + 1)

        # 调整力量
        if keys[pygame.K_DOWN]:
            state.power = max(0, state.power - 5)
        if keys[pygame.K_UP]:
            state.power = min(MAX_POWER, state.power + 5)

        # 发射
        if keys[pygame.K_SPACE]:
            radians = math.radians(state.angle)
            state.ball_vel = [
                state.power * math.cos(radians),
                -state.power * math.sin(radians)
            ]
            state.flying = True
            state.first_impact = True
            state.start_time = time.time()


running = True
clock = pygame.time.Clock()
prev_time = time.time()

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    current_time = time.time()
    dt = current_time - prev_time
    prev_time = current_time

    handle_input()
    update_physics(dt)

    screen.fill(BLACK)
    draw_ground()
    pygame.draw.circle(screen, RED, [int(state.ball_pos[0]),
                                     int(state.ball_pos[1])], BALL_RADIUS)
    if not state.flying and not state.rolling:
        draw_arrow()
    draw_ui()

    pygame.display.flip()
    clock.tick(60)

pygame.quit()